import torch
import torch.nn as nn
import torch.nn.functional as F

class ContrastiveLoss(nn.Module):
	def __call__(self, ypred, label):
		B, _ = ypred.shape
		ypred = F.normalize(ypred)
		cos_matrix = ypred.mm(ypred.t())
		pos_label_matrix = torch.stack([label == label[i] for i in range(B)]).float()
		neg_label_matrix = 1 - pos_label_matrix
		pos_cos_matrix = 1 - cos_matrix
		neg_cos_matrix = cos_matrix - 0.4
		neg_cos_matrix[neg_cos_matrix < 0] = 0
		loss = (pos_cos_matrix * pos_label_matrix).sum() + (neg_cos_matrix * neg_label_matrix).sum()
		loss /= (B * B)
		return loss